"""
The LossSoftmaxCCE class implements combined softmax activation and categorical cross entropy loss.
"""

import numpy as np

class LossSoftmaxCCE:

    def __init__(self):
        """
        No initialization necessary.
        """
        pass

    def forward_pass(self, inputs, y_true):
        """
        Compute softmax activation and cross entropy loss.
        inputs: logits, shape (classes, batch_size)
        y_true: one-hot true labels, shape (classes, batch_size)
        """
        self.y_true = y_true
        # Softmax activation
        exp_values = np.exp(inputs - np.max(inputs, axis=0, keepdims=True))
        self.probs = exp_values / np.sum(exp_values, axis=0, keepdims=True)
        # Cross-entropy loss (negative log likelihood)
        sample_losses = -np.log(self.probs[y_true == 1])
        return np.mean(sample_losses)

    def backward_pass(self):
        """
        Compute gradient of loss with respect to inputs (logits).
        """
        batch_size = self.y_true.shape[1]
        self.dinputs = (self.probs - self.y_true) / batch_size
        return self.dinputs